#!/bin/bash

# --------------------------------------------------------------------------- #
# Pre-defined the necessary constants
# --------------------------------------------------------------------------- #
# Job limit attributes
DYNAMIC_JOB_WALLTIME_LIMIT=36000 # 10 hours
VIRTUAL_MEMORY_LIMIT=100000000   # 100GB
JOB_DEFAULT_WALLTIME=360000 # 100 hours is given if the attribute 'walltime' lost
# Big mem job should request above the largest memory size of static slot
BIG_MEM_JOB_THRESHOLD=4700       # 4.7GB     
# Condor Worker client path
CONDOR_WORKER_WARRPER_CONF="/var/lib/CondorWorker/CondorWorker/wrapper/script.d"
# Job walltime limit config file
JOB_WALLTIME_FILE="/etc/condor/wrapper/walltime.txt"
JOB_SPECIAL_WALLTIME_FILE="/etc/condor/wrapper/walltime_special.txt"


USER_PATH="$PATH"
USER_LDPATH="$LD_LIBRARY_PATH"
export PATH="/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin:${PATH}"
export LD_LIBRARY_PATH="/lib:/lib64:${LD_LIBRARY_PATH}"


# --------------------------------------------------------------------------- #
# Get and define the common viriables
# --------------------------------------------------------------------------- #
# _CONDOR_JOB_AD saves the classad file path of current running job
if [[ ! -f "$_CONDOR_JOB_AD" ]]; then
	/bin/echo "ERROR: Job classad doesn't exist. Please try to re-submit this job. If the error appeared again, please post feedback on helpdesk.ihep.ac.cn." >&2
fi

# Export _CONDOR_IHEP_JOB_ID, user can get job id with this viriable
CONDOR_IHEP_JOB_ID=`/bin/grep -E '^GlobalJobId ' $_CONDOR_JOB_AD | /bin/awk -F'#' '{print $2}'`
export _CONDOR_IHEP_JOB_ID=$CONDOR_IHEP_JOB_ID

# Export _CONDOR_IHEP_REMOTE_HOST, user can get where the job is running on
CONDOR_IHEP_REMOTE_HOST=`/bin/grep -E '^RemoteHost ' $_CONDOR_JOB_AD | /bin/awk -F '[ "]' '{print $4}'`
export _CONDOR_IHEP_REMOTE_HOST="$CONDOR_IHEP_REMOTE_HOST"

# Export _CONDOR_IHEP_SUBMITTED_TIME, user can get when the job was submitted
CONDOR_IHEP_SUBMISSION_TIME=`/bin/grep -E '^QDate ' $_CONDOR_JOB_AD | /bin/awk '{print $3}'`
CONDOR_IHEP_SUBMISSION_TIME=`/bin/date -d @${CONDOR_IHEP_SUBMISSION_TIME} +"%Y-%m-%d %T"`
export _CONDOR_IHEP_SUBMISSION_TIME="$CONDOR_IHEP_SUBMISSION_TIME"

# Export _CONDOR_SCSP_JOB_INFO, this attribute is dedicated for storage troubleshooting
CONDOR_SCSP_GLOBAL_JOB_ID=`/bin/grep -E '^GlobalJobId ' $_CONDOR_JOB_AD | /bin/awk -F '[="]' '{print $3}'`
CONDOR_SCSP_OWNER=`/bin/grep -E '^Owner ' $_CONDOR_JOB_AD | /bin/awk -F '[="]' '{print $3}'`
CONDOR_SCSP_GROUP=`/bin/grep -E '^HepJob_RealGroup ' $_CONDOR_JOB_AD | /bin/awk -F '[="]' '{print $3}'`
export _OPMC_HTC_JOB_INFO="${CONDOR_SCSP_GLOBAL_JOB_ID}#${CONDOR_SCSP_OWNER}#${CONDOR_SCSP_GROUP}"

# Export OPMC_SITE_NAME, user can get which site the job is running at
OPMC_SITE_NAME=`/bin/grep -E "^SiteName " $_CONDOR_MACHINE_AD |awk -F '[="]' '{print $3}'`
export OPMC_SITE_NAME=${OPMC_SITE_NAME}

# Get the number of cpu cores requested by job 
request_cpus=`/bin/grep -E '^RequestCpus ' $_CONDOR_JOB_AD |/bin/awk '{print $3}'`

# Get the memory size requested by job
request_memory=`/bin/grep -E '^RequestMemory ' $_CONDOR_JOB_AD |/bin/awk '{print $3}'` 

# Get job walltime, which means how long time the job would take on running.
# Tt's only effective on long job currently.
#job_walltime=`/bin/grep -E '^WallTime ' $_CONDOR_JOB_AD |/bin/awk -F'"' '{print $2}'`

# Check if requesting singularity, result 0 indicates requesting SL6 singularity
request_singularity=`/bin/grep -E '^IHEP_RequestContainer ' $_CONDOR_JOB_AD | /bin/awk -F '"' '{print $2}'`
if [[ -n "$request_singularity" ]] && [[ "$request_singularity" != "SL6" && "$request_singularity" != "SL7" && "$request_singularity" != "CentOS7" ]];
then	
	/bin/echo "ERROR: No container '${request_singularity}' under serving." >&2
	exit 19
fi 

# Get the owner,acctgroupuser,acctgroup of this job
owner=`/bin/grep -E '^Owner ' $_CONDOR_JOB_AD | /bin/awk -F'"' '{print $2}'`
acctgroupuser=`/bin/grep -E '^AcctGroupUser ' $_CONDOR_JOB_AD |/bin/awk -F'"' '{print $2}'`
acctgroup=`/bin/grep -E '^AcctGroup ' $_CONDOR_JOB_AD |/bin/awk -F'"' '{print $2}'`
real_group=`/bin/grep -E '^IHEP_RealGroup ' $_CONDOR_JOB_AD |/bin/awk -F '"' '{print $2}'`

# Get mem information
#total_mem=`/usr/bin/free -k |grep Mem | awk '{print $2}'`
# Get cpu information
#total_cpu=`/usr/bin/lscpu |grep "^CPU(s):" | awk '{print $2}'`

# Get OS version of this node
os_version=`/bin/sed -r 's/[^0-9]*([0-9]+).*/\1/' /etc/redhat-release`
# --------------------------------------------------------------------------- #
# Verify user permission
# --------------------------------------------------------------------------- #
ownergroups=($(/usr/bin/groups $owner 2>&1))
unset ownergroups[0] ownergroups[1]
# Add 'short' as a legal group
ownergroups=("${ownergroups[@]}" "short" "test" "long")

if [[ ${ownergroups[@]/$real_group/} == ${ownergroups[@]} ]];
then
	/bin/echo "FATAL: Please set your own group." >&2
	exit 20
fi

if [[ "$owner" != "$acctgroupuser" ]];
then
	/bin/echo "FATAL: Please set your own username." >&2
	exit 21
fi
# --------------------------------------------------------------------------- #
# Limits area
# --------------------------------------------------------------------------- #
# Execute user's job with maximum (JobLimit) wall time
job_walltime=`/bin/grep -E '^HepJob_Walltime ' $_CONDOR_JOB_AD |/bin/awk -F'"' '{print $2}'`
if [[ -z $job_walltime ]];
then
	job_walltime="default"
fi
# Check if the current user exists in the walltime special list
special_group=($(/bin/grep -E "^${real_group}_" ${JOB_SPECIAL_WALLTIME_FILE} |/bin/awk -F'=' '{print $1}'))
if [[ -n "${special_group[0]}" ]] && [[ "${job_walltime}" == "default" ]];
then
	special_user_list=($(/bin/grep -E "^${real_group}_" ${JOB_SPECIAL_WALLTIME_FILE} |/bin/awk -F'[="]' '{print $3}'))
	if [[ "${special_user_list[@]}" =~ "${owner}" ]];
	then
		walltime_limit=`/bin/grep -E "^${special_group[0]}_walltime" ${JOB_SPECIAL_WALLTIME_FILE} |/bin/awk -F'=' '{print $2}'`
	fi
fi
# If the current group is not in special list, retrieve walltime value in job walltime file
if [[ -z $walltime_limit ]];
then
	walltime_limit=`/bin/grep -E ${real_group}_${job_walltime}  ${JOB_WALLTIME_FILE} |/bin/awk -F'=' '{print $2}'`
fi
# If walltime_limit is failed to get a value, set as the default job walltime limit
if [[ -z $walltime_limit ]];
then
        walltime_limit=$JOB_DEFAULT_WALLTIME
fi

# Limit virutal memory site available by each job
ulimit -v $VIRTUAL_MEMORY_LIMIT

# --------------------------------------------------------------------------- #
# Singularity area
# --------------------------------------------------------------------------- #
image_sl5="/cvmfs/container.ihep.ac.cn/singularity/image/SL55/WorkNode55-onlyread-20240729.sif"
image_sl6="/cvmfs/container.ihep.ac.cn/singularity/image/SL69/sl69worknode"
image_sl7="/cvmfs/container.ihep.ac.cn/singularity/image/SL7/sl7worknode20200729.sif"
image_centos7="/cvmfs/container.ihep.ac.cn/singularity/image/CentOS/Centosimagelink"
singularity_path="/usr/bin/apptainer"
singularity_cmd="$singularity_path exec"

#export SINGULARITYENV_PATH=$PATH
export APPTAINERENV_SING_USER_DEFINED_APPEND_PATH=$USER_PATH
export APPTAINERENV_LD_LIBRARY_PATH=$USER_LDPATH

if [[ -n "$SHELL" ]] ; 
then
	export APPTAINER_SHELL=$SHELL
else
	current_user=`/bin/whoami`
	user_shell=`/bin/awk -F ":" '{print $1 ":"$NF}' /etc/passwd|grep -w $current_user|/bin/awk -F ":" '{print $2}'`
	export APPTAINER_SHELL=$user_shell
fi	

# --------------------------------------------------------------------------- #
# Sub-Wrapper in Condor Worker System 
# --------------------------------------------------------------------------- #
# Path of sub wrapper script in Condor Worker System
script_path="$_CONDOR_WORKER_WARRPER_CONF/$real_group"

# --------------------------------------------------------------------------- #
# run job
# --------------------------------------------------------------------------- #
# Recieve the singals
trap stop_user_job 2 3 15
stop_user_job() {
    /usr/bin/kill -TERM $job_pid
}

# Pre-scripts
pre_script="$script_path""/pre.sh" 
if [[ -f "$pre_script" ]]; then
	/bin/bash $pre_script $acctgroupuser $real_group
	exit_status=$?
	if [[ $exit_status -ne 0 ]];
	then
		/bin/echo "ERROR: Your job can not be initialized." >&2
		/bin/echo "Please contact helpdesk@ihep.ac.cn." >&2
		exit 24
	fi
fi

# Skip the job script to the remaining parameters
user_exe=$1
shift


export PATH="${USER_PATH}"
export LD_LIBRARY_PATH="${USER_LDPATH}"

# Check if a singularity is needed and pick a proper image
if [[ -z "$request_singularity" ]]; then
	/usr/bin/timeout $walltime_limit ${user_exe} "$@" &
elif [[ "$request_singularity" == "SL5" ]]; then 
	/usr/bin/timeout $walltime_limit $singularity_cmd ${image_sl5} ${user_exe} "$@" &
elif [[ "$request_singularity" == "SL6" ]]; then 
	/usr/bin/timeout $walltime_limit $singularity_cmd ${image_sl6} ${user_exe} "$@" &
elif [[ "$request_singularity" == "SL7" ]]; then
	/usr/bin/timeout $walltime_limit $singularity_cmd ${image_sl7} ${user_exe} "$@" &
elif [[ "$request_singularity" == "CentOS7" ]]; then
	/usr/bin/timeout $walltime_limit $singularity_cmd ${image_centos7} ${user_exe} "$@" &
else
	/usr/bin/timeout $walltime_limit ${user_exe} "$@" &
fi


# Capture the pid of the job to kill later, if needed
job_pid=$!
# Wait for the job to exit of own accord
wait $job_pid
exit_code=$?
if [ $exit_code -eq 124 ]; then
    if [ $walltime_limit -ge 3600 ]; then
        time_limit_hour=`/bin/awk -v walltime_limit=$walltime_limit 'BEGIN{printf "%.1f\n",walltime_limit/3600}'`
        /bin/echo "Job is KILLED: exceeds the ${time_limit_hour} hour(s) time limit" >&2
    elif [ $walltime_limit -lt 3600 ]; then
        time_limit_minute=`/bin/awk -v walltime_limit=$walltime_limit 'BEGIN{printf "%.1f\n",walltime_limit/60}'`
        /bin/echo "Job is KILLED: exceeds the ${time_limit_minute} minute(s) time limit" >&2
    fi
fi

# When the job exitted, put code to close out
post_script="$script_path/post.sh" 
if [[ -f "$post_script" ]]; then
	/bin/bash $post_script $acctgroupuser $real_group
	exit_status=$?
	if [[ $exit_status -ne 0 ]]; then
		/bin/echo "ERROR: Your job can not be finalized" >&2
		/bin/echo "Please contact helpdesk@ihep.ac.cn." >&2
		exit 29
	fi
fi

# Exit with the jobs exit code
exit $exit_code
